package com.ics.atable.chat.controller;

import com.alibaba.cloud.ai.graph.*;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.checkpoint.constant.SaverEnum;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.StateSnapshot;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.ics.atable.chat.controller.graph.GraphProcess;
import com.ics.atable.chat.expection.ErrorCode;
import com.ics.atable.chat.expection.ThrowUtils;
import com.ics.atable.chat.model.dto.ChatRequest;
import com.ics.atable.chat.model.dto.FeedbackRequest;
import com.ics.atable.chat.model.vo.node.stream.ObjectStreamResponse;
import com.ics.atable.chat.model.vo.node.stream.StreamResponse;
import com.ics.atable.chat.model.entity.ChatMessage;
import com.ics.atable.chat.model.enums.ResponseTypeEnum;
import com.ics.atable.chat.service.ChatMessageService;
import com.ics.atable.chat.serializer.ChatContentCodec;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static com.ics.atable.chat.constant.ItemQueryConstant.*;

/**
 * 指标查询控制层
 *
 * @author haohongbin
 */
@Slf4j
@RestController
@RequestMapping("/item-query")
public class ItemQueryController {

    private final CompiledGraph compiledGraph;

    private final ChatMessageService chatMessageService;

    // 使用本地缓存存储 session 到 Sink 的映射
    private final Map<String, Sinks.Many<ServerSentEvent<StreamResponse>>> sessionSinkMap = new ConcurrentHashMap<>();

    private final ChatContentCodec chatContentCodec;

    public ItemQueryController(@Qualifier("itemQueryGraph") StateGraph stateGraph, ChatMessageService chatMessageService,
                               ChatContentCodec chatContentCodec)
            throws GraphStateException {
//        SaverConfig saverConfig = SaverConfig.builder().register(SaverEnum.REDIS.getValue(), new RedisSaver()).build();
        SaverConfig saverConfig = SaverConfig.builder().register(SaverEnum.MEMORY.getValue(), new MemorySaver()).build();
        this.compiledGraph = stateGraph.compile(CompileConfig.builder().saverConfig(saverConfig).interruptBefore(HUMAN_FEEDBACK).build());
        this.chatMessageService = chatMessageService;
        this.chatContentCodec = chatContentCodec;
    }

    @RequestMapping(value = "/chat/stream", method = RequestMethod.POST, produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<StreamResponse>> chatStream(@RequestBody(required = false) ChatRequest chatRequest) throws GraphRunnerException {
        ThrowUtils.throwIf(ObjectUtils.isEmpty(chatRequest), ErrorCode.PARAMS_ERROR, "请求参数不能为空");
        ThrowUtils.throwIf(ObjectUtils.isEmpty(chatRequest.sessionId()), ErrorCode.PARAMS_ERROR, "会话ID不能为空");
        ThrowUtils.throwIf(ObjectUtils.isEmpty(chatRequest.type()), ErrorCode.PARAMS_ERROR, "会话类型不能为空");
        ThrowUtils.throwIf(ObjectUtils.isEmpty(ResponseTypeEnum.getByCode(chatRequest.type())), ErrorCode.PARAMS_ERROR, "会话类型不合法");
        ThrowUtils.throwIf(ObjectUtils.isEmpty(chatRequest.query()) && ObjectUtils.isEmpty(chatRequest.interruptFeedback()),
                ErrorCode.PARAMS_ERROR, "用户输入或中断反馈不能为空");

        // 存储输入
        ChatMessage chatMessage = new ChatMessage();
        chatMessage.setSessionId(chatRequest.sessionId());
        chatMessage.setRole(MessageType.USER.getValue());
        chatMessage.setContent(chatRequest.query());
        chatMessage.setMessageType(chatRequest.type());
        chatMessageService.save(chatMessage);

        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(chatRequest.sessionId()).build();

        Map<String, Object> objectMap = new HashMap<>();
        // Create a unicast sink to emit ServerSentEvents
        Sinks.Many<ServerSentEvent<StreamResponse>> sink = Sinks.many().unicast().onBackpressureBuffer();
        sessionSinkMap.put(chatRequest.sessionId(), sink);

        GraphProcess graphProcess = new GraphProcess(this.compiledGraph);
        // Handle human feedback if auto-accept is disabled and feedback is provided
        if (StringUtils.hasText(chatRequest.interruptFeedback())) {
            graphProcess.handleHumanFeedback(chatRequest, objectMap, runnableConfig, sink);
        }
        // First question
        else {
            objectMap.put("query", chatRequest.query());
            objectMap.put("session_id", chatRequest.sessionId());
            log.info("init inputs: {}", objectMap);
            AsyncGenerator<NodeOutput> resultFuture = compiledGraph.stream(objectMap, runnableConfig);
            graphProcess.processStream(resultFuture, sink, ResponseTypeEnum.getByCode(chatRequest.type()));
        }

        // 收集所有发送的数据，结束后保存到数据库
        StringBuilder fullResponse = new StringBuilder();

        ChatMessage response = new ChatMessage();
        response.setSessionId(chatRequest.sessionId());
        response.setRole(MessageType.ASSISTANT.getValue());
        response.setMessageType(chatRequest.type());

        List<ObjectStreamResponse> responseList = new ArrayList<>();

        return sink.asFlux()
                .doOnNext(serverSentEvent -> {
                    // 将每个响应片段添加到完整响应中
                    if (serverSentEvent.data() != null) {
                        fullResponse.append(serverSentEvent.data().toString());
                        // 如果是ObjectStreamResponse类型，也单独收集
                        if (serverSentEvent.data() instanceof ObjectStreamResponse) {
                            responseList.add((ObjectStreamResponse) serverSentEvent.data());
                        }
                    }
                })
                .doOnComplete(() -> {
                    // 流结束时将完整响应保存到数据库
                    String jsonContent = chatContentCodec.serializeObjectStreamResponses(responseList);
                    response.setContent(jsonContent);
                    chatMessageService.save(response);
                })
                .doOnCancel(() -> {
                    log.info("Client disconnected from stream");
                    // 客户端断开连接时也保存已收集的数据
                    response.setContent(fullResponse.toString());
                    chatMessageService.save(response);
                })
                .doOnError(e -> {
                    log.error("Error occurred during streaming", e);
                    // 出错时保存已收集的数据
                    response.setContent(fullResponse.toString());
                    chatMessageService.save(response);
                });
    }

    @RequestMapping(value = "/chat/resume", method = RequestMethod.POST, produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<StreamResponse>> resume(@RequestBody(required = false) FeedbackRequest humanFeedback) throws GraphRunnerException {
        ThrowUtils.throwIf(ObjectUtils.isEmpty(humanFeedback), ErrorCode.PARAMS_ERROR, "参数不能为空");
        ThrowUtils.throwIf(ObjectUtils.isEmpty(humanFeedback.sessionId()), ErrorCode.PARAMS_ERROR, "sessionId不能为空");
        ThrowUtils.throwIf(humanFeedback.feedBack() && ObjectUtils.isEmpty(humanFeedback.feedBackContent()), ErrorCode.PARAMS_ERROR, "反馈不能为空");
        ThrowUtils.throwIf(ObjectUtils.isEmpty(humanFeedback.type()), ErrorCode.PARAMS_ERROR, "会话类型不能为空");
        ThrowUtils.throwIf(ObjectUtils.isEmpty(ResponseTypeEnum.getByCode(humanFeedback.type())), ErrorCode.PARAMS_ERROR, "会话类型不合法");

        // 存储输入
        ChatMessage chatMessage = new ChatMessage();
        chatMessage.setSessionId(humanFeedback.sessionId());
        chatMessage.setRole(MessageType.USER.getValue());
        chatMessage.setContent(humanFeedback.feedBackContent());
        chatMessage.setMessageType(humanFeedback.type());
        chatMessageService.save(chatMessage);

        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(humanFeedback.sessionId()).build();
        Map<String, Object> objectMap = new HashMap<>();
        objectMap.put("feedback", humanFeedback.feedBack());
        objectMap.put("feedback_content", humanFeedback.feedBackContent());

        // Create a unicast sink to emit ServerSentEvents
        Sinks.Many<ServerSentEvent<StreamResponse>> sink = Sinks.many().unicast().onBackpressureBuffer();
        GraphProcess graphProcess = new GraphProcess(this.compiledGraph);

        StateSnapshot stateSnapshot = compiledGraph.getState(runnableConfig);
        OverAllState state = stateSnapshot.state();
        state.withResume();
        state.withHumanFeedback(new OverAllState.HumanFeedback(objectMap, INTENT));

        AsyncGenerator<NodeOutput> resultFuture = compiledGraph.streamFromInitialNode(state, runnableConfig);
        graphProcess.processStream(resultFuture, sink, ResponseTypeEnum.getByCode(humanFeedback.type()));

        // 收集所有发送的数据，结束后保存到数据库
        StringBuilder fullResponse = new StringBuilder();

        ChatMessage response = new ChatMessage();
        response.setSessionId(humanFeedback.sessionId());
        response.setRole(MessageType.ASSISTANT.getValue());
        response.setMessageType(humanFeedback.type());

        return sink.asFlux()
                .doOnNext(serverSentEvent -> {
                    // 将每个响应片段添加到完整响应中
                    if (serverSentEvent.data() != null) {
                        fullResponse.append(serverSentEvent.data());
                    }
                })
                .doOnComplete(() -> {
                    // 流结束时将完整响应保存到数据库
                    response.setContent(fullResponse.toString());
                    chatMessageService.save(response);
                })
                .doOnCancel(() -> {
                    log.info("Client disconnected from stream");
                    // 客户端断开连接时也保存已收集的数据
                    response.setContent(fullResponse.toString());
                    chatMessageService.save(response);
                })
                .doOnError(e -> {
                    log.error("Error occurred during streaming", e);
                    // 出错时保存已收集的数据
                    response.setContent(fullResponse.toString());
                    chatMessageService.save(response);
                });
    }

    @RequestMapping(value = "/chat/close", method = RequestMethod.GET)
    public void closeStream(@RequestParam(value = "sessionId") String sessionId) {
        // 由于每个请求都有自己的 sink，这里不再需要关闭全局 sink
        // 如果需要实现关闭功能，需要维护一个 session 到 sink 的映射
        Sinks.Many<ServerSentEvent<StreamResponse>> sink = sessionSinkMap.get(sessionId);
        if (sink != null) {
            sink.tryEmitComplete();
            sessionSinkMap.remove(sessionId);
            log.info("Closed stream for session: {}", sessionId);
        } else {
            log.warn("No active stream found for session: {}", sessionId);
        }
    }
}
